#code to replicate analysis of Lalonde data
rm(list=ls())
library(foreign)
library(Matching)
library(survey)
library(lattice)
# package that implements entropy balancing
library(ebal)

dat <- read.table("cps1re74.csv",header=T)
# unemployed
dat$u74 <- as.numeric(dat$re74==0)
dat$u75 <- as.numeric(dat$re75==0)

# covars
covars <- c("age","educ","black","hispan","married","nodegree","re74","re75","u74","u75")

# compute all interactions
X <- as.matrix(dat[,covars])
XX <- matrixmaker(X)

# prepare data: exclude non-sensical and co-linear vars, code treatment and outcome, change names
out <- c("black.black","nodegree.nodegree","nodegree.educ","married.married","hispan.hispan","hispan.black",
         "u75.u75","u74.u74","u74.re74","u75.re75","u74.re74","u75.re75","re74.re74","re75.re75","re75.re74")

XX     <- XX[,(colnames(XX) %in% out)==F]
dat    <- data.frame(dat[,(names(dat) %in% covars)==F],XX)
covars <- names(dat)[-which((names(dat) %in% c("treat","re78")))]
X     <- dat[,covars]
X     <- as.matrix(X)
colnames(X) <-gsub(".age","*Age",colnames(X))
colnames(X) <-gsub(".educ*","*Schooling",colnames(X))
colnames(X) <-gsub(".black","*Black",colnames(X))
colnames(X) <-gsub(".hispan","*Hispanic",colnames(X))
colnames(X) <-gsub(".married","*Married",colnames(X))
colnames(X) <-gsub(".re74","*Earnings 1974",colnames(X))
colnames(X) <-gsub(".re75","*Earnings 1975",colnames(X))
colnames(X) <-gsub(".u74","*Unemployed 1974",colnames(X))
colnames(X) <-gsub(".u75","*Unemployed 1975",colnames(X))
colnames(X) <-gsub(".nodegree","*HS Dropout",colnames(X))
colnames(X) <-gsub("age","Age",colnames(X))
colnames(X) <-gsub("educ","Schooling",colnames(X))
colnames(X) <-gsub("black","Black",colnames(X))
colnames(X) <-gsub("hispan","Hispanic",colnames(X))
colnames(X) <-gsub("married","Married",colnames(X))
colnames(X) <-gsub("re74","Earnings 1974",colnames(X))
colnames(X) <-gsub("re75","Earnings 1975",colnames(X))
colnames(X) <-gsub("u74","Unemployed 1974",colnames(X))
colnames(X) <-gsub("u75","Unemployed 1975",colnames(X))
colnames(X) <-gsub("nodegree","HS Dropout",colnames(X))
dat <- data.frame(dat,X)
Y <- dat$re78
W <- dat$tr

#  balance before matching
bout.nm   <- MatchBalance(W~X,match.out = NULL,ks=FALSE)
bal.nm <- baltest.collect(matchbal.out=bout.nm ,var.names=colnames(X),after=FALSE)
round(bal.nm,3)

# Maha Dist Matching
mout.maha  <-  Match(Y,W,X,BiasAdjust=F,estimand="ATT",M=1)
summary(mout.maha)
bout.maha  <- MatchBalance(W~X,match.out = mout.maha,ks=FALSE)
bal.maha <- baltest.collect(matchbal.out=bout.maha ,var.names=colnames(X),after=TRUE)
round(bal.maha,3)

# Genetic Matching ATT
g.weights <- GenMatch(Tr=W, X=X, BalanceMatrix=X, estimand="ATT", M=1,print.level=0)
mout.gm <-  Match(Y,W,X,BiasAdjust=F,Weight.matrix=g.weights,estimand="ATT",M=1)
summary(mout.gm)
bout.gm <- MatchBalance(W~X,match.out = mout.gm,print.level=0,ks=FALSE)
bal.gm <- baltest.collect(matchbal.out=bout.gm,var.names=colnames(X),after=TRUE)
round(bal.gm,3)

## Logistic PS weighting + matching
PS  <- glm(W~X,family=binomial(link=logit))$fitted
PSM <- PS
PS  <- PS[W==0]
PS  <- PS/(1-PS)

# PS Matching 
mout.psm <-  Match(Y,W,X=PSM,BiasAdjust=F,estimand="ATT",M=1)
summary(mout.psm)
bout.psm <- MatchBalance(W~X,match.out = mout.psm,print.level=0,ks=FALSE)
bal.psm  <- baltest.collect(matchbal.out=bout.psm,var.names=colnames(X),after=TRUE)
round(bal.psm,3)

# PS Weighting
bout.psw <- MatchBalance(W~X,weights=c(W[W==1],PS),ks=FALSE)
bal.psw <- baltest.collect(matchbal.out=bout.psw,var.names=colnames(X),after=FALSE)
round(bal.psw,2)

# Entropy Balancing
out.eb <- ebalance(
              Treatment=W,
              X=X
              )
           
bout.eb <- MatchBalance(W~X,weights=c(W[W==1],out.eb$w),ks=FALSE)
bal.eb  <- baltest.collect(matchbal.out=bout.eb,var.names=colnames(X),after=F)
round(bal.eb,2)

# Entropy Balancing (with trimmed weights)
out.ebtr <- ebalance.trim(          
               ebalanceobj=out.eb,
              )

bout.ebtr <- MatchBalance(W~X,weights=c(W[W==1],out.ebtr$w),ks=FALSE)
bal.ebtr  <- baltest.collect(matchbal.out=bout.eb,var.names=colnames(X),after=FALSE)
round(bal.ebtr,2)

dat$w   <-  c(W[W==1],out.eb$w)
dat$wtr <-  c(W[W==1],out.ebtr$w)

# Effect Estimates
des1 <- svydesign(id=~1,weights=~w, data= dat)
mod1 <- svyglm(re78 ~ treat, design = des1)
summary(mod1)

des2 <- svydesign(id=~1,weights=~wtr, data= dat)
mod2 <- svyglm(re78 ~ treat, design = des2)
summary(mod2)

# Balance Plots
d  <- rbind(bal.nm,bal.maha,bal.gm,bal.psm,bal.psw,bal.eb)
dn <- rownames(d)
d  <- data.frame(d)
d$vname <- factor(dn,levels = unique(dn)[length(unique(dn)):1], labels = unique(dn)[length(unique(dn)):1])
d$gr <- rep(c("Unadjusted","MahaDist Matching","Genetic Matching","PS Matching","PS Weighting","Entropy Balancing"),each=length(unique(dn)))
d$gr <- factor(d$gr,levels=unique(d$gr)[1:length(unique(d$gr))],labels=unique(d$gr)[1:length(unique(d$gr))])

mypal<-c("black",rep("darkgrey",4),"black")[6:1]
Cex <- Cex2 <- 1

# plot with SDs
bplot <- function(x,y,...)
 {
 panel.abline(v=0, lwd = 1 , lty="solid")
 panel.abline(v=c(-.1,.1), lwd = 2 , lty="dotted")
 panel.abline(h=c(1:nrow(d)), lwd = 1 , lty="dashed", col="gray95")
 panel.xyplot(x,y,...)
 }

print(

xyplot(vname~round(sdiff.pooled/100,2),data=d,groups=gr,xlim=c(-.6,.6),
panel = bplot
 ,par.settings = list(superpose.symbol = list(pch = c(15,19,17,18,15,1)[6:1],col=mypal,cex=1))
 ,xlab=list("standardized difference in means",cex=Cex),ylab="",auto.key=T,scales=list(y=list(cex=Cex),x=list(cex=Cex2,at=c(-.5,-.1,0,.1,.5),labels=c("-.5","-.1","0",".1",".5")))
 )
)
# plot with p-values
print(
xyplot(vname~round(T.pval,2),data=d,groups=gr,xlim=c(-.05,1.05),
panel = bplot
 ,par.settings = list(superpose.symbol = list(pch = c(15,19,17,18,15,1)[6:1],col=mypal,cex=1))
 ,xlab=list("p-value: difference of means test",cex=Cex),ylab="",auto.key=T,scales=list(y=list(cex=Cex),x=list(cex=Cex2,at=c(0,.1,.2,.5,1),labels=c("0",".1",".2",".5","1")))
 )
)

# table
latex(round(cbind(bal.nm[,c(1:2)],bal.nm[,4]/100,bal.nm[,5:6],
                  bal.eb[,2],bal.eb[,4]/100,bal.eb[,5:6],
                  bal.psw[,2],bal.psw[,4]/100,bal.psw[,5:6]),2),file="")


# Sensitivity Simulation
covarstosample <- covars 
sims  <- 1000000
store <- matrix(NA,sims,2)
colnames(store) <- c("EB","OLS")
for(i in 1:sims){
 print(i) 
        ftmp <- "re78 ~ treat"
        bk   <- sample(covarstosample,sample(1:length(covarstosample),1),replace=F)
        ftmp <- as.formula(paste(ftmp, "+",paste(bk,collapse=" + ")))
        modw       <- svyglm(ftmp, design = des1)
        modols <- lm(ftmp, data = dat) 
 store[i,1:2] <-  c(coef(modw)["treat"],coef(modols)["treat"])
}

# Estimates  
summary(store)
dx1 <- density(store[,2])
# Plot
plot(dx1, type = "l", ylab = "Density", 
         xlab = "Estimated average treatment effect for the treated (US$)",
         xlim = range(c(dx1$x)), 
         ylim = range(c(dx1$y)),main="")
segments(y0=0,y1=0.001232,x1=mean(store[,"EB"]),x0=mean(store[,"EB"]),lty="dashed",lwd=2)
arrows(y0=.001,x0=000,y1=0.001,x1=1600,length=.1) 
text(y=.001,x=-2600,labels="Estimates after Entropy Balancing") 
arrows(y0=.0002,x0=-2000,y1=0.0002,x1=0,length=.1) 
text(y=.0002,x=-4800,labels="Estimates before Entropy Balancing")

# QQ plots for distributions
des <- svydesign(id=~1,weights=~w, data= dat[dat$treat==1,])
tt  <- svyquantile(~re75+re74+age+educ, des, seq(0,1,by=.01), ci=TRUE)$quantiles
des <- svydesign(id=~1,weights=~w, data= dat[dat$treat==0,])
cc  <- svyquantile(~re75+re74+age+educ, des, seq(0,1,by=.01), ci=TRUE)$quantiles

par(mfrow=c(2,2))
nnames <- c("Earnings 1975","Earnings 1974","Age","Schooling")
for(i in 1:4){
qqplot(dat[dat$treat==1,rownames(tt)[i]],dat[dat$treat==0,rownames(tt)[i]],
      main=paste("QQ plot of ",nnames[i],sep=""),
      xlab=c("treated units"),ylab=c("control units"),
      pch=19
      )
Pum <- qqplot(y=tt[i,],x=cc[i,], plot.it = F)
points(Pum$x,Pum$y,pch=19,col="gray")
abline(coef=c(0,1),col="red")
}

